package com.pragmatic.cryptography;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * RSA加解密demo
 *
 * @author wangdong
 * @since 2025/4/21 16:53
 */
public class SimpleRsaCryptography {

  // 算法名称
  public static final String KEY_ALGORITHM = "RSA";
  // 定义签名算法
  private static final String KEY_RSA_SIGNATURE = "MD5withRSA";
  // 公钥和私钥生成后一般配置起来(这里demo暂存到map，这里由于是模拟发送方和接受方，所以定义两对密钥对)
  private static final String SEND_PUBLIC_KEY = "send_rsa_public_key";
  private static final String SEND_PRIVATE_KEY = "send_rsa_private_key";

  private static final String RECEIVE_PUBLIC_KEY = "receive_rsa_public_key";
  private static final String RECEIVE_PRIVATE_KEY = "receive_rsa_private_key";

  // ==============================================================================
  // ===========================base64转换字符串和字节数组=============================
  // ==============================================================================

  /**
   * BASE64 解码 （Str -> byte[]）
   *
   * @param key 需要Base64解码的字符串
   * @return 字节数组
   */
  public static byte[] decryptBase64(String key) {
    return Base64.getDecoder().decode(key);
  }

  /**
   * BASE64 编码 （byte[] -> Str）
   *
   * @param key 需要Base64编码的字节数组
   * @return 字符串
   */
  public static String encryptBase64(byte[] key) {
    return new String(Base64.getEncoder().encode(key));
  }

  /**
   * 获取Base64编码的公钥/私钥字符串
   *
   * @param map 默认map
   * @param paramKey 公私钥key
   * @return 秘钥String
   */
  public static String getSecurityKey(Map<String, Object> map, String paramKey) {
    String str = "";
    Key key = (Key) map.get(paramKey);
    str = encryptBase64(key.getEncoded());
    return str;
  }

  // ==============================================================================
  // =================================生成密钥对=====================================
  // ==============================================================================

  /**
   * 初始化公钥、私钥对
   *
   * @return map
   * @throws Exception 生成过程异常
   */
  public static Map<String, Object> initKey() throws Exception {
    KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(KEY_ALGORITHM);
    keyPairGen.initialize(1024);
    KeyPair keyPair = keyPairGen.generateKeyPair();
    // 发送方：公钥和私钥对
    RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
    RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
    Map<String, Object> keyMap = new HashMap<>(2);
    keyMap.put(SEND_PUBLIC_KEY, publicKey);
    keyMap.put(SEND_PRIVATE_KEY, privateKey);

    KeyPairGenerator keyPairGen2 = KeyPairGenerator.getInstance(KEY_ALGORITHM);
    keyPairGen2.initialize(1024);
    KeyPair keyPair2 = keyPairGen2.generateKeyPair();
    // 接收方：公钥和私钥对
    RSAPublicKey publicKey2 = (RSAPublicKey) keyPair2.getPublic();
    RSAPrivateKey privateKey2 = (RSAPrivateKey) keyPair2.getPrivate();
    keyMap.put(RECEIVE_PUBLIC_KEY, publicKey2);
    keyMap.put(RECEIVE_PRIVATE_KEY, privateKey2);

    return keyMap;
  }

  // ==============================================================================
  // ==============================发送方：对内容加解密================================
  // ==============================================================================

  /**
   * 公钥加密(对方公钥加密)
   *
   * @param contentStr 待加密的报文
   * @param publicKeyStr 对方公钥
   * @return 加密后的密文
   */
  public static String encryptByPublic(String contentStr, String publicKeyStr) throws Exception {
    // 获得公钥
    byte[] publicKeyBytes = decryptBase64(publicKeyStr);
    // 将公钥由字符串转为UTF-8格式的字节数组
    X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyBytes);
    // 取得待加密数据
    byte[] data = contentStr.getBytes(StandardCharsets.UTF_8);
    KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
    PublicKey publicKey = factory.generatePublic(keySpec);
    // 对数据加密
    Cipher cipher = Cipher.getInstance(factory.getAlgorithm());
    cipher.init(Cipher.ENCRYPT_MODE, publicKey);
    // 返回加密后由Base64编码的加密信息
    return encryptBase64(cipher.doFinal(data));
  }

  /**
   * 私钥解密(己方私钥解密)
   *
   * @param encryptedStr 被己方公钥加密的密文
   * @param privateKeyStr 己方私钥
   * @return 明文
   */
  public static String decryptByPrivate(String encryptedStr, String privateKeyStr)
      throws Exception {
    // 对私钥解密
    byte[] privateKeyBytes = decryptBase64(privateKeyStr);
    // 获得私钥
    PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
    // 获得待解密数据
    byte[] data = decryptBase64(encryptedStr);
    KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
    PrivateKey privateKey = factory.generatePrivate(keySpec);
    // 对数据解密
    Cipher cipher = Cipher.getInstance(factory.getAlgorithm());
    cipher.init(Cipher.DECRYPT_MODE, privateKey);
    // 返回UTF-8编码的解密信息
    return new String(cipher.doFinal(data), StandardCharsets.UTF_8);
  }

  // ==============================================================================
  // =================================加签和验签===================================
  // ==============================================================================

  /**
   * 用私钥对加密数据进行签名（己方私钥对数据进行签名）
   *
   * @param encryptedStr 待签名密文(被对方公钥加密的数据)
   * @param privateKey 己方私钥
   * @return 签名后的密文
   */
  public static String sign(String encryptedStr, String privateKey) throws Exception {
    // 将私钥加密数据字符串转换为字节数组
    byte[] data = encryptedStr.getBytes();
    // 解密由base64编码的私钥
    byte[] bytes = decryptBase64(privateKey);
    // 构造PKCS8EncodedKeySpec对象
    PKCS8EncodedKeySpec pkcs = new PKCS8EncodedKeySpec(bytes);
    // 指定的加密算法
    KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
    // 取私钥对象
    PrivateKey key = factory.generatePrivate(pkcs);
    // 用私钥对信息生成数字签名
    Signature signature = Signature.getInstance(KEY_RSA_SIGNATURE);
    signature.initSign(key);
    signature.update(data);
    return encryptBase64(signature.sign());
  }

  /**
   * 校验数字签名 (对方用公钥验签)
   *
   * @param encryptedStr 待验签密文
   * @param publicKey 对方公钥
   * @param sign 签名数据
   * @return 校验成功返回true，失败返回false
   */
  public static boolean verify(String encryptedStr, String publicKey, String sign)
      throws Exception {
    // 将私钥加密数据字符串转换为字节数组
    byte[] data = encryptedStr.getBytes();
    // 解密由base64编码的公钥
    byte[] bytes = decryptBase64(publicKey);
    // 构造X509EncodedKeySpec对象
    X509EncodedKeySpec keySpec = new X509EncodedKeySpec(bytes);
    // 指定的加密算法
    KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
    // 取公钥对象
    PublicKey key = factory.generatePublic(keySpec);
    // 用公钥验证数字签名
    Signature signature = Signature.getInstance(KEY_RSA_SIGNATURE);
    signature.initVerify(key);
    signature.update(data);
    return signature.verify(decryptBase64(sign));
  }

  // main方法测试
  public static void main(String[] args) throws Exception {
    // 1、生成密钥对
    Map<String, Object> keyMap = SimpleRsaCryptography.initKey();
    // 2、发送方A首先通过接收方B的公钥加密内容
    String content = "你好吗？";
    System.out.println("原始数据===" + content);
    String encryptedContent =
        SimpleRsaCryptography.encryptByPublic(content, getSecurityKey(keyMap, RECEIVE_PUBLIC_KEY));
    // 3、发送方A再通过己方的私钥对内容签名
    String sign =
        SimpleRsaCryptography.sign(encryptedContent, getSecurityKey(keyMap, SEND_PRIVATE_KEY));
    // 4、发送数据Data(2产生) + sign(3产生)
    // 这里的data和sign如果被截获，针对data没有私钥没办法打开(防止泄露)；如果三方用B的公钥篡改了data,因为没有A的私钥所以没法改签名数据(改了sign没法签名或者签名肯定不对)，
    // B在接收到数据验签发现签名的报文和三方篡改的报文不一致验签不通过
    System.out.println("加密数据===" + encryptedContent);
    System.out.println("签名数据===" + sign);
    // 5、接收方B用发送方A的公钥验证签名
    boolean verify =
        SimpleRsaCryptography.verify(
            encryptedContent, getSecurityKey(keyMap, SEND_PUBLIC_KEY), sign);
    System.out.println("验签结果===" + verify);
    // 6、接收方B用己方私钥对A发送的密文解密
    String decryptContent =
        SimpleRsaCryptography.decryptByPrivate(
            encryptedContent, getSecurityKey(keyMap, RECEIVE_PRIVATE_KEY));
    System.out.println("解密数据===" + decryptContent);
  }
}
